import os
from typing import Tuple
import numpy as np
from tqdm import tqdm
from loadDataset import LoadBCICompDataSet


def MakeDataSetAndSave(matFilePath: str, saveRootPath: str) -> None:
    '''
    Description:
        这个函数是用来将mat文件中是Target和不是Target的信号进行分类，并保存到文件中

        文件的顺序如下
            └─Subject
                └─NoTargetSignal
                    └─NoTarget_Num
                └─TargetSignal
                    └─Target_Num


    Parameters:
        matFilePath:str
            mat文件所在的路径
        saveRootPath:str
            TargetNum所在的根文件夹
    '''
    dataSet = LoadBCICompDataSet.LoadBCICompDataSet(matFilePath)
    signal: np.array = dataSet["Signal"]
    flash: np.array = dataSet["Flashing"]
    code: np.array = dataSet["StimulusCode"]
    stimulateType: np.array = dataSet["StimulusType"]
    saveTarData = np.zeros((85, 64, 12, 15, 240))
    saveNoTarData = np.zeros((85, 64, 12, 15, 240))
    windows = 240
    for i in tqdm(range(0, np.size(signal, 0)), desc="Making..."):
        testSignal: np.array = signal[i]
        testFlash: np.array = flash[i]
        testCode: np.array = code[i]
        testStimulusType: np.array = stimulateType[i]
        for j in range(0, np.size(testSignal, 1)):
            # print(test[:, i].size)
            currentChannel = testSignal[:, j]
            # print(type(currentChannel.size)) # int
            rowColCount = np.zeros(12)
            for k in range(1, testFlash.size):
                if testFlash[k] == 0 and testFlash[k - 1] == 1:
                    if testStimulusType[k] == 0 and testStimulusType[k - 1] == 1:
                        rowCol = int(testCode[k - 1]) - 1
                        saveTarData[i, j, rowCol, int(rowColCount[rowCol]), :] = currentChannel[
                                                                                 k - 24:k + windows - 24]  # 按照案例给的
                    else:
                        rowCol = int(testCode[k - 1]) - 1
                        saveNoTarData[i, j, rowCol, int(rowColCount[rowCol]), :] = currentChannel[
                                                                                   k - 24:k + windows - 24]  # 按照案例给的
    root = saveRootPath
    if not os.path.exists(root):
        os.makedirs(root)
    os.chdir(root)
    targetSignalPath = os.path.join(root, "TargetSignal")
    if not os.path.exists(targetSignalPath):
        os.makedirs(targetSignalPath)
    os.chdir(targetSignalPath)
    for i in tqdm(range(0, np.size(saveTarData, 0))):
        targetPath = os.path.join(targetSignalPath, "Target_" + str(i + 1))
        if not os.path.exists(targetPath):
            os.mkdir(targetPath)
        os.chdir(targetPath)
        for j in range(0, np.size(saveTarData, 1)):
            a = saveTarData[i, j, :, :]
            # print(a.shape)
            np.save("channel_" + str(j + 1) + ".npy", a)

    noTargetSignalPath = os.path.join(root, "NoTargetSignal")
    if not os.path.exists(noTargetSignalPath):
        os.makedirs(noTargetSignalPath)
    os.chdir(noTargetSignalPath)
    for i in tqdm(range(0, np.size(saveNoTarData, 0))):
        targetPath = os.path.join(noTargetSignalPath, "Target_" + str(i + 1))
        if not os.path.exists(targetPath):
            os.mkdir(targetPath)
        os.chdir(targetPath)
        for j in range(0, np.size(saveNoTarData, 1)):
            a = saveNoTarData[i, j, :, :]
            # print(a.shape)
            np.save("channel_" + str(j + 1) + ".npy", a)


def MakeDataSet(matFilePath: str) -> Tuple[np.ndarray, np.ndarray]:
    '''
    Description:
        这个函数是用来将mat文件中是Target和不是Target的信号进行分类，不会保存到文件中
        注意使用的时候需要使用numpy.all排除全是0的数组

    Parameters:
        matFilePath:str
            mat文件所在的路径

    Returns:
        Tuple[np.ndarray, np.ndarray]

            返回saveTarData, saveNoTarData

            保存Target和NoTarget的信号数组，大小都是[85, 64, 12, 15, 240]
    '''
    dataSet = LoadBCICompDataSet.LoadBCICompDataSet(matFilePath)
    signal: np.array = dataSet["Signal"]
    flash: np.array = dataSet["Flashing"]
    code: np.array = dataSet["StimulusCode"]
    stimulateType: np.array = dataSet["StimulusType"]
    saveTarData = np.zeros((85, 64, 12, 15, 240))
    saveNoTarData = np.zeros((85, 64, 12, 15, 240))
    windows = 240
    for i in tqdm(range(0, np.size(signal, 0)), desc="Making..."):
        testSignal: np.array = signal[i]
        testFlash: np.array = flash[i]
        testCode: np.array = code[i]
        testStimulusType: np.array = stimulateType[i]
        for j in range(0, np.size(testSignal, 1)):
            # print(test[:, i].size)
            currentChannel = testSignal[:, j]
            # print(type(currentChannel.size)) # int
            rowColCount = np.zeros(12)
            for k in range(1, testFlash.size):
                if testFlash[k] == 0 and testFlash[k - 1] == 1:
                    if testStimulusType[k] == 0 and testStimulusType[k - 1] == 1:
                        rowCol = int(testCode[k - 1]) - 1
                        saveTarData[i, j, rowCol, int(rowColCount[rowCol]), :] = currentChannel[
                                                                                 k - 24:k + windows - 24]  # 按照案例给的
                    else:
                        rowCol = int(testCode[k - 1]) - 1
                        saveNoTarData[i, j, rowCol, int(rowColCount[rowCol]), :] = currentChannel[
                                                                                   k - 24:k + windows - 24]  # 按照案例给的
    return saveTarData, saveNoTarData


def LoadTarAndNoTarData(subjectRootPath: str, convertToF: bool = True) -> Tuple[np.ndarray, np.ndarray]:
    '''
    Description:
        这个函数是用来读取分类保存到文件中的numpy数组
        注意使用的时候需要使用numpy.all排除全是0的数组

    Parameters:
        subjectRootPath:str
            分类后的根路径
        convertToF: bool
            是否转换为Float32类型

    Returns:
        Tuple[np.ndarray, np.ndarray]

            返回saveTarData, saveNoTarData

            保存Target和NoTarget的信号数组，大小都是[85, 64, 12, 15, 240]
    '''
    tarDatas = np.zeros((85, 64, 12, 15, 240))
    noTarDatas = np.zeros((85, 64, 12, 15, 240))
    tarPath = os.path.join(subjectRootPath, "TargetSignal")
    noTarPath = os.path.join(subjectRootPath, "NoTargetSignal")
    paths = os.listdir(tarPath)
    countTarget = 0
    for i in tqdm(paths, desc='Loading {} Target data'.format(os.path.basename(subjectRootPath))):
        path = os.path.join(tarPath, i)
        files = os.listdir(path)
        countChannel = 0
        for j in files:
            file = os.path.join(path, j)
            data = np.load(file)
            tarDatas[countTarget, countChannel, :, :, :] = data
            countChannel += 1
        countTarget += 1

    paths = os.listdir(noTarPath)
    countTarget = 0
    for i in tqdm(paths, desc='Loading {} NoTarget data'.format(os.path.basename(subjectRootPath))):
        path = os.path.join(noTarPath, i)
        files = os.listdir(path)
        countChannel = 0
        for j in files:
            file = os.path.join(path, j)
            data = np.load(file)
            noTarDatas[countTarget, countChannel, :, :, :] = data
            countChannel += 1
        countTarget += 1
    if convertToF:
        tarDatas = tarDatas.astype(np.float32)
        noTarDatas = noTarDatas.astype(np.float32)
    return tarDatas, noTarDatas


if __name__ == "__main__":
    MakeDataSetAndSave(r"F:\DataSet\BCICompDIII\DataSet\BCI_Comp_III_Wads_2004\Subject_A_Train.mat",
                       r"F:\DataSet\BCICompDIII\DataSet\MyDataSet\SubjectA")
    MakeDataSetAndSave(r"F:\DataSet\BCICompDIII\DataSet\BCI_Comp_III_Wads_2004\Subject_B_Train.mat",
                       r"F:\DataSet\BCICompDIII\DataSet\MyDataSet\SubjectB")
    # a, b = LoadTarAndNoTarData(r"F:\DataSet\BCICompDIII\DataSet\MyDataSet\SubjectA")
    # print(a[0, 0, 0, 0, :])
